package connection

import (
	"github.com/panjf2000/gnet/v2"
	"sync"
	"sync/atomic"
)

type ConnMatrix interface {
	AddConn(c gnet.Conn)
	LoadCount() (n int32)
	GetConn(fd int) gnet.Conn
	DelConn(c gnet.Conn)
}
type connMatrix struct {
	connCount int32
	connMap   map[int]gnet.Conn
	mx        sync.Mutex
}

func NewConnMatrix() ConnMatrix {
	cm := &connMatrix{}
	cm.init()
	return cm
}

func (cm *connMatrix) init() {
	cm.connMap = make(map[int]gnet.Conn)
}

func (cm *connMatrix) iterate(f func(gnet.Conn) bool) {
	for _, c := range cm.connMap {
		if c != nil {
			if !f(c) {
				return
			}
		}
	}
}

func (cm *connMatrix) incCount(_ int, delta int32) {
	atomic.AddInt32(&cm.connCount, delta)
}

func (cm *connMatrix) LoadCount() (n int32) {
	return atomic.LoadInt32(&cm.connCount)
}

func (cm *connMatrix) AddConn(c gnet.Conn) {
	cm.connMap[c.Fd()] = c
	cm.incCount(0, 1)
}

func (cm *connMatrix) DelConn(c gnet.Conn) {
	delete(cm.connMap, c.Fd())
	cm.incCount(0, -1)
}

func (cm *connMatrix) GetConn(fd int) gnet.Conn {
	return cm.connMap[fd]
}
